import json
import os
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.patches import FancyBboxPatch, ConnectionPatch
import numpy as np
from collections import defaultdict, deque
from typing import Dict, List, Any, Tuple, Optional
import math

class HierarchicalTreeVisualizer:
    """
    专门的层次化树可视化器
    生成真正的树状结构图
    """
    
    def __init__(self):
        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
        plt.rcParams['axes.unicode_minus'] = False
        
        # 节点样式配置
        self.node_styles = {
            'FILE': {
                'color': '#2E7D32', 'text_color': 'white', 'width': 150, 'height': 50,
                'emoji': '📄', 'font_size': 11, 'border_width': 3
            },
            'FUNCTION_DEF': {
                'color': '#1976D2', 'text_color': 'white', 'width': 140, 'height': 45,
                'emoji': '🔧', 'font_size': 10, 'border_width': 2
            },
            'ANONYMOUS_FUNCTION': {
                'color': '#FF9800', 'text_color': 'white', 'width': 130, 'height': 40,
                'emoji': '🎯', 'font_size': 9, 'border_width': 2
            },
            'CLASS': {
                'color': '#9C27B0', 'text_color': 'white', 'width': 140, 'height': 45,
                'emoji': '🏛️', 'font_size': 10, 'border_width': 2
            },
            'CALL': {
                'color': '#F57C00', 'text_color': 'white', 'width': 120, 'height': 35,
                'emoji': '📞', 'font_size': 9, 'border_width': 2
            },
            'ARGUMENT': {
                'color': '#7B1FA2', 'text_color': 'white', 'width': 100, 'height': 30,
                'emoji': '📋', 'font_size': 8, 'border_width': 1
            },
            'PARAMETER': {
                'color': '#607D8B', 'text_color': 'white', 'width': 90, 'height': 28,
                'emoji': '📝', 'font_size': 8, 'border_width': 1
            }
        }
        
        # 边样式配置
        self.edge_styles = {
            'defines': {'color': '#3F51B5', 'width': 3.0, 'alpha': 0.9, 'style': '-'},
            'calls': {'color': '#FF5722', 'width': 2.5, 'alpha': 0.8, 'style': '-'},
            'has_arg': {'color': '#9C27B0', 'width': 1.8, 'alpha': 0.7, 'style': '--'},
            'has_param': {'color': '#607D8B', 'width': 1.5, 'alpha': 0.6, 'style': ':'},
            'nested_in': {'color': '#FF9800', 'width': 2.0, 'alpha': 0.8, 'style': '-.'},
            'chained_call': {'color': '#E91E63', 'width': 2.5, 'alpha': 0.9, 'style': '-'},
            'callback': {'color': '#795548', 'width': 1.8, 'alpha': 0.7, 'style': '--'}
        }
        
        # 布局配置
        self.layout_config = {
            'level_height': 120,        # 层间距离
            'min_node_spacing': 40,     # 最小节点间距
            'subtree_spacing': 60,      # 子树间额外间距
            'margin_x': 150,           # 左右边距
            'margin_y': 80,            # 上下边距
            'curve_strength': 0.3      # 连接线弯曲强度
        }
    
    def load_skeleton_graph(self, json_file_path: str) -> Dict[str, Any]:
        """加载骨架图数据"""
        with open(json_file_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    
    def build_tree_structure(self, graph_data: Dict) -> Tuple[Dict, str, Dict]:
        """构建树结构"""
        # 构建邻接表
        children = defaultdict(list)
        parents = {}
        all_nodes = {node['id']: node for node in graph_data['nodes']}
        
        # 处理边关系
        for edge in graph_data['edges']:
            parent_id = edge['source']
            child_id = edge['target']
            edge_type = edge['type']
            
            if parent_id in all_nodes and child_id in all_nodes:
                children[parent_id].append({
                    'id': child_id,
                    'edge_type': edge_type
                })
                parents[child_id] = {
                    'id': parent_id,
                    'edge_type': edge_type
                }
        
        # 找到根节点
        root_node = None
        for node in graph_data['nodes']:
            if node['id'] not in parents and node['type'] == 'FILE':
                root_node = node['id']
                break
        
        if root_node is None:
            # 如果没有FILE类型的根节点，选择第一个没有父节点的节点
            for node in graph_data['nodes']:
                if node['id'] not in parents:
                    root_node = node['id']
                    break
        
        if root_node is None and graph_data['nodes']:
            # 最后的备选：选择第一个节点
            root_node = graph_data['nodes'][0]['id']
        
        return children, root_node, all_nodes
    
    def calculate_tree_layout(self, children: Dict, root_node: str, all_nodes: Dict) -> Dict[str, Tuple[float, float]]:
        """计算完美的树布局"""
        if not root_node or root_node not in all_nodes:
            return {}
        
        # 第一步：计算每个节点的层级
        levels = self._calculate_levels(children, root_node)
        
        # 第二步：计算每个子树的宽度
        subtree_widths = self._calculate_subtree_widths(children, root_node, all_nodes)
        
        # 第三步：计算水平位置
        positions = self._calculate_positions(children, root_node, levels, subtree_widths, all_nodes)
        
        return positions
    
    def _calculate_levels(self, children: Dict, root_node: str) -> Dict[str, int]:
        """计算每个节点的层级"""
        levels = {}
        
        def dfs(node_id: str, level: int):
            levels[node_id] = level
            for child_info in children.get(node_id, []):
                dfs(child_info['id'], level + 1)
        
        dfs(root_node, 0)
        return levels
    
    def _calculate_subtree_widths(self, children: Dict, root_node: str, all_nodes: Dict) -> Dict[str, float]:
        """计算每个子树的宽度"""
        widths = {}
        
        def calculate_width(node_id: str) -> float:
            if node_id in widths:
                return widths[node_id]
            
            node_children = children.get(node_id, [])
            
            if not node_children:
                # 叶子节点
                node_style = self.node_styles.get(all_nodes[node_id]['type'], self.node_styles['CALL'])
                widths[node_id] = node_style['width'] + self.layout_config['min_node_spacing']
            else:
                # 内部节点：子树宽度之和
                child_widths = [calculate_width(child['id']) for child in node_children]
                total_child_width = sum(child_widths)
                
                # 子树间需要额外间距
                if len(node_children) > 1:
                    total_child_width += (len(node_children) - 1) * self.layout_config['subtree_spacing']
                
                # 节点自身宽度
                node_style = self.node_styles.get(all_nodes[node_id]['type'], self.node_styles['CALL'])
                node_width = node_style['width'] + self.layout_config['min_node_spacing']
                
                widths[node_id] = max(total_child_width, node_width)
            
            return widths[node_id]
        
        calculate_width(root_node)
        return widths
    
    def _calculate_positions(self, children: Dict, root_node: str, levels: Dict, 
                           subtree_widths: Dict, all_nodes: Dict) -> Dict[str, Tuple[float, float]]:
        """计算节点位置"""
        positions = {}
        
        def position_subtree(node_id: str, x_center: float):
            # 设置当前节点位置
            level = levels[node_id]
            y = -level * self.layout_config['level_height']  # 从上到下
            positions[node_id] = (x_center, y)
            
            # 处理子节点
            node_children = children.get(node_id, [])
            if not node_children:
                return
            
            # 计算子节点的起始位置
            total_width = 0
            for child_info in node_children:
                child_id = child_info['id']
                total_width += subtree_widths[child_id]
            
            # 子树间间距
            if len(node_children) > 1:
                total_width += (len(node_children) - 1) * self.layout_config['subtree_spacing']
            
            # 子节点的起始x坐标
            start_x = x_center - total_width / 2
            current_x = start_x
            
            # 递归定位每个子树
            for child_info in node_children:
                child_id = child_info['id']
                child_width = subtree_widths[child_id]
                child_center_x = current_x + child_width / 2
                
                position_subtree(child_id, child_center_x)
                
                current_x += child_width + self.layout_config['subtree_spacing']
        
        # 从根节点开始定位
        position_subtree(root_node, 0)
        
        return positions
    
    def generate_tree_visualization(self, graph_data: Dict, output_path: str, 
                                  figsize: Tuple[int, int] = (30, 20)):
        """生成树状可视化"""
        print("🌳 开始生成层次化树状图...")
        
        # 1. 构建树结构
        children, root_node, all_nodes = self.build_tree_structure(graph_data)
        
        if not root_node:
            print("❌ 未找到根节点，无法生成树状图")
            return
        
        print(f"🌱 根节点: {root_node}")
        print(f"📊 节点统计: {len(all_nodes)} 个节点")
        
        # 2. 计算布局
        positions = self.calculate_tree_layout(children, root_node, all_nodes)
        
        if not positions:
            print("❌ 布局计算失败")
            return
        
        print(f"📐 布局计算完成: {len(positions)} 个节点有位置")
        
        # 3. 创建图形
        fig, ax = plt.subplots(figsize=figsize, dpi=100)
        ax.set_aspect('equal')
        
        # 4. 绘制连接线
        self._draw_tree_connections(children, positions, all_nodes, ax)
        
        # 5. 绘制节点
        self._draw_tree_nodes(positions, all_nodes, ax)
        
        # 6. 设置标题和样式
        self._setup_plot_appearance(ax, graph_data)
        
        # 7. 调整视图范围
        self._adjust_view_limits(positions, ax)
        
        # 8. 保存
        self._save_plot(fig, output_path)
        
        print(f"✅ 树状图生成完成: {output_path}")
    
    def _draw_tree_connections(self, children: Dict, positions: Dict, all_nodes: Dict, ax):
        """绘制树的连接线"""
        print("🔗 绘制连接线...")
        
        drawn_connections = 0
        
        for parent_id, child_list in children.items():
            if parent_id not in positions:
                continue
            
            parent_x, parent_y = positions[parent_id]
            
            for child_info in child_list:
                child_id = child_info['id']
                edge_type = child_info['edge_type']
                
                if child_id not in positions:
                    continue
                
                child_x, child_y = positions[child_id]
                
                # 获取边样式
                style = self.edge_styles.get(edge_type, {
                    'color': '#666666', 'width': 2.0, 'alpha': 0.7, 'style': '-'
                })
                
                # 绘制弯曲连接线
                self._draw_curved_connection(
                    parent_x, parent_y, child_x, child_y, 
                    style, ax
                )
                
                drawn_connections += 1
        
        print(f"   绘制了 {drawn_connections} 条连接线")
    
    def _draw_curved_connection(self, x1: float, y1: float, x2: float, y2: float, 
                              style: Dict, ax):
        """绘制弯曲的连接线"""
        # 计算控制点
        mid_y = (y1 + y2) / 2
        curve_strength = self.layout_config['curve_strength']
        
        # 创建贝塞尔曲线
        t = np.linspace(0, 1, 50)
        
        # 三次贝塞尔曲线控制点
        p0 = np.array([x1, y1])
        p1 = np.array([x1, y1 - abs(y2 - y1) * curve_strength])
        p2 = np.array([x2, y2 + abs(y2 - y1) * curve_strength])
        p3 = np.array([x2, y2])
        
        # 计算曲线点
        curve_points = np.array([
            (1-t)**3 * p0[i] + 3*(1-t)**2*t * p1[i] + 3*(1-t)*t**2 * p2[i] + t**3 * p3[i]
            for i in range(2)
        ]).T
        
        # 绘制曲线
        ax.plot(curve_points[:, 0], curve_points[:, 1],
               color=style['color'],
               linewidth=style['width'],
               alpha=style['alpha'],
               linestyle=style['style'],
               zorder=1)
        
        # 添加箭头
        # 箭头位置在曲线的80%处
        arrow_idx = int(len(curve_points) * 0.8)
        arrow_start_idx = max(0, arrow_idx - 3)
        
        if arrow_idx < len(curve_points) and arrow_start_idx < len(curve_points):
            ax.annotate('', 
                       xy=curve_points[arrow_idx], 
                       xytext=curve_points[arrow_start_idx],
                       arrowprops=dict(
                           arrowstyle='->', 
                           color=style['color'],
                           lw=style['width'] * 0.8,
                           alpha=style['alpha']
                       ),
                       zorder=2)
    
    def _draw_tree_nodes(self, positions: Dict, all_nodes: Dict, ax):
        """绘制树节点"""
        print("⭕ 绘制节点...")
        
        drawn_nodes = 0
        
        for node_id, (x, y) in positions.items():
            if node_id not in all_nodes:
                continue
            
            node = all_nodes[node_id]
            node_type = node['type']
            style = self.node_styles.get(node_type, self.node_styles['CALL'])
            
            # 创建节点标签
            label = self._create_node_label(node)
            
            # 绘制节点背景
            width = style['width']
            height = style['height']
            
            # 圆角矩形
            bbox = FancyBboxPatch(
                (x - width/2, y - height/2),
                width, height,
                boxstyle="round,pad=5",
                facecolor=style['color'],
                edgecolor='white',
                linewidth=style['border_width'],
                alpha=0.9,
                zorder=3
            )
            ax.add_patch(bbox)
            
            # 添加文本
            ax.text(x, y, label,
                   ha='center', va='center',
                   fontsize=style['font_size'],
                   color=style['text_color'],
                   weight='bold',
                   zorder=4,
                   wrap=True)
            
            drawn_nodes += 1
        
        print(f"   绘制了 {drawn_nodes} 个节点")
    
    def _create_node_label(self, node: Dict) -> str:
        """创建节点标签"""
        node_type = node['type']
        style = self.node_styles.get(node_type, {})
        emoji = style.get('emoji', '❓')
        
        if node_type == 'FILE':
            return f"{emoji}\nJavaScript File"
        
        elif node_type == 'FUNCTION_DEF':
            func_name = node.get('function_name', 'unknown')
            if len(func_name) > 15:
                func_name = func_name[:15] + "..."
            return f"{emoji}\n{func_name}()"
        
        elif node_type == 'ANONYMOUS_FUNCTION':
            is_callback = node.get('is_callback', False)
            if is_callback:
                return f"{emoji}\n(callback)"
            else:
                return f"{emoji}\n(anonymous)"
        
        elif node_type == 'CLASS':
            class_name = node.get('class_name', 'unknown')
            if len(class_name) > 12:
                class_name = class_name[:12] + "..."
            return f"{emoji}\nclass {class_name}"
        
        elif node_type == 'CALL':
            callee = node.get('callee_name', 'unknown')
            # 处理长函数名
            if len(callee) > 18:
                parts = callee.split('.')
                if len(parts) > 1:
                    # 显示最后一部分
                    callee = f"...{parts[-1]}"
                else:
                    callee = callee[:15] + "..."
            
            # 检查是否是链式调用
            if node.get('is_chained', False):
                return f"{emoji}\n{callee}\n(chained)"
            else:
                return f"{emoji}\n{callee}()"
        
        elif node_type == 'ARGUMENT':
            arg_index = node.get('argument_index', 0)
            arg_type = node.get('argument_type', 'unknown')
            
            # 简化参数类型显示
            type_map = {
                'StringLiteral': 'str',
                'NumericLiteral': 'num',
                'BooleanLiteral': 'bool',
                'ObjectExpression': 'obj',
                'ArrayExpression': 'arr',
                'Identifier': 'var',
                'CallExpression': 'call',
                'FunctionExpression': 'func',
                'TemplateLiteral': 'tmpl',
                'RegExpLiteral': 'regex'
            }
            
            short_type = type_map.get(arg_type, arg_type[:4])
            return f"{emoji}\narg{arg_index}\n<{short_type}>"
        
        elif node_type == 'PARAMETER':
            param_name = node.get('parameter_name', 'param')
            if len(param_name) > 10:
                param_name = param_name[:10] + "..."
            return f"{emoji}\n{param_name}"
        
        else:
            return f"{emoji}\n{node_type}"
    
    def _setup_plot_appearance(self, ax, graph_data: Dict):
        """设置图形外观"""
        # 设置标题
        ax.set_title("JavaScript 代码调用树结构", 
                    fontsize=18, fontweight='bold', pad=30)
        
        # 添加图例
        legend_elements = []
        for node_type, style in self.node_styles.items():
            emoji = style.get('emoji', '❓')
            display_name = node_type.replace('_', ' ').title()
            legend_elements.append(
                mpatches.Patch(color=style['color'], 
                              label=f"{emoji} {display_name}")
            )
        
        # 分两列显示图例
        legend1 = ax.legend(handles=legend_elements[:4], 
                           title="节点类型", 
                           loc='upper left',
                           bbox_to_anchor=(0.02, 0.98),
                           fontsize=9,
                           title_fontsize=10)
        
        if len(legend_elements) > 4:
            legend2 = ax.legend(handles=legend_elements[4:], 
                               title="", 
                               loc='upper left',
                               bbox_to_anchor=(0.02, 0.78),
                               fontsize=9)
            ax.add_artist(legend1)
        
        # 添加边类型说明
        edge_legend_elements = []
        for edge_type, style in self.edge_styles.items():
            edge_legend_elements.append(
                plt.Line2D([0], [0], color=style['color'], linewidth=style['width'],
                          label=edge_type.replace('_', ' ').title())
            )
        
        ax.legend(handles=edge_legend_elements, 
                 title="连接类型", 
                 loc='upper right',
                 bbox_to_anchor=(0.98, 0.98),
                 fontsize=8,
                 title_fontsize=9)
        
        ax.set_axis_off()
    
    def _adjust_view_limits(self, positions: Dict, ax):
        """调整视图范围"""
        if not positions:
            return
        
        x_coords = [pos[0] for pos in positions.values()]
        y_coords = [pos[1] for pos in positions.values()]
        
        x_min, x_max = min(x_coords), max(x_coords)
        y_min, y_max = min(y_coords), max(y_coords)
        
        # 添加边距
        margin_x = self.layout_config['margin_x']
        margin_y = self.layout_config['margin_y']
        
        ax.set_xlim(x_min - margin_x, x_max + margin_x)
        ax.set_ylim(y_min - margin_y, y_max + margin_y)
    
    def _save_plot(self, fig, output_path: str):
        """保存图形"""
        try:
            plt.tight_layout()
            fig.savefig(output_path, dpi=300, bbox_inches='tight',
                       facecolor='white', edgecolor='none',
                       pad_inches=0.3)
            print(f"💾 图形已保存: {output_path}")
        except Exception as e:
            print(f"❌ 保存失败: {e}")
        finally:
            plt.close(fig)


def main():
    """主函数"""
    # 输入输出路径
    input_file = r"E:\PDCG\malicious\1_postmessage-api-external-3.25.12\package\build.js.pdcg.json"
    output_dir = r"d:\PythonProject\malicious_npm_detector_1.0.1\struct_analysis"
    
    if not os.path.exists(input_file):
        print(f"❌ 输入文件不存在: {input_file}")
        return
    
    print("🌳 启动层次化树状图生成器...")
    print(f"📁 输入文件: {input_file}")
    
    # 创建可视化器
    visualizer = HierarchicalTreeVisualizer()
    
    try:
        # 加载数据
        print("📊 加载骨架图数据...")
        graph_data = visualizer.load_skeleton_graph(input_file)
        
        print(f"📊 数据统计：{len(graph_data.get('nodes', []))} 节点，{len(graph_data.get('edges', []))} 边")
        
        # 验证数据
        if graph_data.get('validation', {}).get('is_valid_tree', False):
            print("✅ 数据验证通过：有效的树结构")
        else:
            print("⚠️ 数据验证警告：可能不是完美的树结构")
        
        # 生成输出文件名
        base_name = os.path.splitext(os.path.basename(input_file))[0]
        output_path = os.path.join(output_dir, f"{base_name}_perfect_tree.png")
        
        # 生成层次化树状图
        visualizer.generate_tree_visualization(graph_data, output_path, (36, 24))
        
        print("✅ 层次化树状图生成完成!")
        print(f"📁 输出文件: {output_path}")
        
    except Exception as e:
        print(f"❌ 生成失败: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()